-- 本文件是 init_by_lua 阶段加载
-- 所有以 g_ 开头的都是全局变量,其他地方直接使用即可

g_config = require 'config'
g_util = require 'lib.util'

g_rules_path = g_config.cls_waf_base_path.."/rules"

g_rule_ip_white = g_util.load_rule('ip_white')
g_rule_ip_black = g_util.load_rule('ip_black')
g_rule_user_agent = g_util.load_rule('user_agent')
g_rule_user_agent_white = g_util.load_rule('user_agent_white')
g_rule_url = g_util.load_rule('url')
g_rule_url_white = g_util.load_rule('url_white')
g_rule_args = g_util.load_rule('args')
g_rule_cookie = g_util.load_rule('cookie')
g_rule_post = g_util.load_rule('post')

-- waf日志文件handle
g_log_file_handle = io.open(g_config.waf_log_file, "a")
if not g_log_file_handle then
    ngx.log(ngx.ERR, ' cls_lua_waf log file open failed at init parse.')
end

local match_fun = ngx.re.find

-- IP CIDR工具类
g_ipmatcher = nil
if (require "ffi").os == "Windows" then
    g_ipmatcher = require("lib.ipmatcher_win")
else
    g_ipmatcher = require("lib.ipmatcher")
end

-- 来访IP在白名单内返回true, 未开启白名单检测功能/来访IP不在名单内,返回false
function ip_white_check(client_ip)
    if g_config.ip_white_check == 'on' then
        local matcher = g_ipmatcher.new(g_rule_ip_white)
        if matcher then
            local ok, err = matcher:match(client_ip)
            if ok == true then
                if g_config.ip_white_log == 'on' then
                    g_util.waf_log('[rule:ip_white]')
                end
                return true
            end
        end
    end
    return false
end

-- 来访IP在黑名单内返回true, 未开启黑名单检测功能/来访IP不在名单内,返回false
function ip_black_check(client_ip)
    if g_config.ip_black_check == 'on' then
        local matcher = g_ipmatcher.new(g_rule_ip_black)
        if matcher then
            local ok, err = matcher:match(client_ip)
            if ok == true then
                g_util.waf_log('[rule:ip_black]')
                return true
            end
        end
    end
    return false
end

-- 如果开启了 境外IP禁止访问,加载libmaxminddb以及Country ip库
if g_config.ip_foreign_check == 'on' then
    -- g_geo 全局变量,其他地方只管用
    g_geo = require 'lib.maxminddb'
    if not g_geo.initted() then
        g_geo.init(g_config.cls_waf_base_path.."/data/Country.mmdb")
    end
end
-- 来访IP是境外IP返回true, 未开启境外IP检测功能/来访IP非境外IP,返回false
function ip_foreign_check(client_ip)
    if g_config.ip_foreign_check == 'on' then
        --support ipv6 e.g. 2001:4860:0:1001::3004:ef68
        local res,err = g_geo.lookup(client_ip)
        if not res then
            ngx.log(ngx.ERR,'failed to LOOKUP ip:['..client_ip..'] reason:',err)
            return false
        end
        -- res -> table  {"country":{"iso_code":"CN"}}
        if res.country and res.country.iso_code then
            local country_code = res.country.iso_code
            if country_code == "CN" or country_code == "HK" or country_code == "MO" or country_code == "TW" then
                -- 如果是 CN/HK/MO/TW 则是国内IP
                return false
            else
                g_util.waf_log('[rule:ip_foreign] "'..client_ip..'"')
                return true
            end
        end
    end
    return false
end

-- 来访UA在UA白名单(搜索引擎蜘蛛要放行)内返回true, 来访UA不在规则名单内,返回false
function user_agent_white_check(ua)
    if g_config.user_agent_white_check == 'on' then
        for _, rule in pairs(g_rule_user_agent_white) do
            local from = match_fun(ua, rule, 'ijos')
            if from then
                if g_config.user_agent_white_log == 'on' then
                    g_util.waf_log('[rule:ua_white] "'..rule..'"')
                end
                return true
            end
        end
    end
    return false
end

-- 来访UA在规则名单内返回true, 未开启UA检测功能/来访UA不在规则名单内,返回false
function user_agent_check(ua)
    if g_config.user_agent_check == 'on' then
        for _, rule in pairs(g_rule_user_agent) do
            local from = match_fun(ua, rule, 'ijos')
            if from then
                g_util.waf_log('[rule:ua] "'..rule..'"')
                return true
            end
        end
    end
    return false
end

-- 全局变量
g_cc_count = tonumber(string.match(g_config.cc_rate,'(.*)/'))
g_cc_duration = tonumber(string.match(g_config.cc_rate,'/(.*)'))
g_cctoken_prefix = 'cls_waf:cctoken:'
g_ccip_prefix = 'cls_waf:ccip:'
-- CC检测可能需要用到redis,而 resty.redis 不能在 init_by_lua 阶段使用(详情见https://github.com/openresty/lua-resty-redis#limitations),所以CC检测方法放在 access_by_lua 阶段
-- function cc_check(ip)
-- end

-- URL在规则白名单内返回true, 未开启URL白名单检测功能/来访URL不在规则名单内,返回false
function url_white_check(uri)
    if g_config.url_white_check == 'on' then
        for _, rule in pairs(g_rule_url_white) do
            local from = match_fun(uri, rule, 'ijos')
            if from then
                if g_config.url_white_log == 'on' then
                    g_util.waf_log('[rule:url_white] "'..rule..'"')
                end
                return true
            end
        end
    end
    return false
end

-- URL在规则内返回true, 未开启URL检测功能/来访URL不在规则名单内,返回false
function url_check(uri)
    if g_config.url_check == 'on' then
        for _, rule in pairs(g_rule_url) do
            local from = match_fun(uri, rule, 'ijos')
            if from then
                g_util.waf_log('[rule:uri] "'..rule..'"')
                return true
            end
        end
    end
    return false
end

-- args在规则内返回true, 未开启args检测功能/args不在规则名单内,返回false
function args_check()
    if g_config.args_check == 'on' then
        local args = ngx.req.get_uri_args()
        for _, rule in pairs(g_rule_args) do
            for name, value in pairs(args) do
                local data = nil
                if type(value) == 'table' then
                    data = table.concat(value, ' ')
                else
                    data = value
                end

                local from = match_fun(ngx.unescape_uri(data), rule, "ijos")
                if data and type(data) ~= "boolean" and rule ~="" and from then
                    g_util.waf_log('[rule:args] "'..rule..'"')
                    return true
                end
            end
        end
    end
    return false
end

function cookie_check()
    if g_config.cookie_check == 'on' then
        local cookie = ngx.var.http_cookie
        for _, rule in pairs(g_rule_cookie) do
            local from = match_fun(cookie, rule, 'ijos')
            if from then
                g_util.waf_log('[rule:cookie] "'..rule..'"')
                return true
            end
        end
    end
    return false
end

function post_check()
    local method = ngx.req.get_method()
    if method == 'POST' or method == 'PUT' then
        ngx.req.read_body()
        local post_data = ngx.req.get_post_args()
        if not post_data then
            return false
        end

        for _, rule in pairs(g_rule_post) do
            for name, value in pairs(post_data) do
                local data = nil
                if type(value) == 'table' then
                    data = table.concat(value, ', ')
                else
                    data = value
                end

                local from = match_fun(ngx.unescape_uri(data), rule, "ijos")
                if data and type(data) ~= "boolean" and data ~="" and from then
                    g_util.waf_log('[rule:post] "'..rule..'"')
                    return true
                end
            end
        end
    end
    return false
end
